import numpy as np
import traffic_snapshot_pb2 as ts

## Given a filename of a protobuf file, reads in a traffic matrix time series and 
## returns the timeseries of traffic matrix snapshots
## (optional) - considered_network_id : Tells the algorithm to consider only traffic from within
## 				just the set of considered_network_ids. If it is None, then the algorithm will just read in everything
##				Defaults to None
def read_traffic_matrix_protobuf(protobuf_filename, considered_network_id=None, return_timestamps=False):
	traffic_matrix_snapshots = []
	traffic_snapshots = ts.TrafficFlowTimeseries()
	try:
		f = open(protobuf_filename, "r")
		traffic_snapshots.ParseFromString(f.read())
		f.close()
	except IOError:
		print("Could not open file: {}".format(protobuf_filename))
	pod_id_set = set() ## first pass is to find out all the pod IDs
	if considered_network_id is None:
		for traffic_snapshot in traffic_snapshots.snapshots:
			for traffic_flow in traffic_snapshot.flows:
				src_id = traffic_flow.src
				dst_id = traffic_flow.dst
				pod_id_set.add(src_id)
				pod_id_set.add(dst_id)
	else:
		pod_id_set = set(considered_network_id)
	pod_id = 0
	pod_id_to_original_id_map = {}
	original_id_to_pod_id_map = {}
	for original_id in pod_id_set:
		pod_id_to_original_id_map[pod_id] = original_id
		original_id_to_pod_id_map[original_id] = pod_id
		pod_id += 1
	number_of_pods = len(pod_id_set)
	print("There number of pods is : {}".format(number_of_pods))
	## now go through another pass to finally form the traffic matrices
	timestamps = [0] * len(traffic_snapshots.snapshots)
	for traffic_snapshot, time_index in zip(traffic_snapshots.snapshots, range(len(traffic_snapshots.snapshots))):
		traffic_matrix = np.zeros((number_of_pods, number_of_pods,))
		unix_ts = traffic_snapshot.timestamp
		timestamps[time_index] = unix_ts
		for traffic_flow in traffic_snapshot.flows:
			src_original_id = traffic_flow.src
			dst_original_id = traffic_flow.dst
			if src_original_id in pod_id_set and dst_original_id in pod_id_set:
				if src_original_id != dst_original_id:
					src_pod_id = original_id_to_pod_id_map[src_original_id]
					dst_pod_id = original_id_to_pod_id_map[dst_original_id]
					traffic_matrix[src_pod_id][dst_pod_id] = traffic_flow.size
		traffic_matrix_snapshots.append(traffic_matrix)
	if return_timestamps:
		return traffic_matrix_snapshots, timestamps
	else:
		return traffic_matrix_snapshots

def read_in_network_ids(filename):
	address = {}
	with open(filename, "r") as f:
		for line in f:
			line_str = line.strip()
			#line_str = line_str.rstrip("L")
			current_address = long(line_str)
			address[current_address] = True
	return address.keys()

## Normalizes the traffic matrix to norm
def normalize_traffic_matrix(traffic_matrix, norm=1.):
	n = len(traffic_matrix)
	normalized_tm = np.zeros((n, n))
	total_sum = sum([sum(row) for row in traffic_matrix])
	for i in range(n):
		for j in range(n):
			normalized_tm[i][j] = traffic_matrix[i][j] / total_sum * norm 
	return normalized_tm